# 层级结构的脚本
import numpy as np
# 绘制进度条
from tqdm import tqdm
import statsmodels.api as sm
import Regression


# 获取层级结构
def Get_HS(data_input):
    # 获取行数,列数
    row_len, col_len = data_input.shape
    # print(row_len)
    # print(col_len)
    # 获取每一行的数据
    levels = []
    for i in tqdm(range(row_len)):
        level = 1
        for j in range(len(data_input.iloc[i, :])):
            if data_input.iloc[i, j] == data_input.iloc[i, j]:
                # 循环到不是None的时候赋值
                level = j + 1
                break
        levels.append(level)
    # print('levels:\n', levels)
    return levels


# 获取循环对象的索引,传入等级即可
# 需要dataframe里面有level一列
# 默认最小值为最上层
# 默认最大值为最下层

def Get_IterationObject_Index(data):
    all_iteration_object_index = []
    # 获取最下层和最上层的水表ID值
    MaxID = np.max(data['level'])
    MinID = np.min(data['level'])
    print('寻找迭代对象:')
    for i in tqdm(range(len(data))):
        # 获得等级值
        current_level = data.iloc[i, :]['level']
        iteration_objects = []
        # 是否是可用数据
        if len(data.iloc[i, :]) > 11:
            if current_level == 1:
                # 获取下面要迭代的对象
                iteration_data_Down = data.iloc[i + 1:, :]

                # # 获取上面要迭代的对象
                # iteration_data_Up = data.iloc[:i:, :]

                # 记录当前的索引值
                current_index = i
                # 判断是否不为最下层水表
                if current_level < MaxID:
                    # 获取下层水表
                    for j in range(len(iteration_data_Down.index)):
                        # 获取当前往下循环的水表的水表等级
                        iteration_objects_level = iteration_data_Down.iloc[j, :]['level']

                        if iteration_objects_level > current_level:
                            # 如果找到下层A的下层B分支，删除掉B层的上层A
                            if len(iteration_objects) >= 1:
                                if iteration_objects_level > iteration_data_Down.iloc[j - 1, :]['level']:
                                    # 删除上层A操作
                                    iteration_objects.pop(-1)
                            # 添加下层
                            iteration_objects.append(i + j + 1)
                        elif iteration_objects_level <= current_level:
                            # 如果小于为上层水表,或同层水表
                            break
                # # 判断是否不为最上层水表
                # if current_level > MinID:
                #     # 获取上层水表
                #     # 从当前位置往上找
                #     for j in range(len(iteration_data_Up.index)):
                #         # 获取当前往上循环的水表的水表等级
                #         iteration_objects_index = iteration_data_Up.iloc[-(j + 1), :]['level']
                #         # 如果小于为上层水表
                #         if iteration_objects_index < current_level:
                #             iteration_objects.append(i - (j + 1))
                #             break

        all_iteration_object_index.append(iteration_objects)
    return all_iteration_object_index


# 对每个行进行回归,需要回归对象的索引值列表
# 警告:直接输入回归数据,去掉不需要的列
def Mul_Regression(data, regression_list):
    # 初始化拟合优度
    R_Squares = []
    print('逐行回归：')
    # 正态检验统计量
    Normal_Results = []
    for i in tqdm(range(len(data))):
        # 如果回归对象大于0才回归
        if len(regression_list[i]):
            # 用于回归的字典
            regression_dict = {}

            regression_dict['y'] = np.array(data.iloc[i, :])
            # 用于回归的字符串
            regression_str = 'y~'
            # 循环加入每个回归行,这里区分上层水表和下层水表
            for j in range(len(regression_list[i])):
                obj_index = regression_list[i][j]
                varname = 'x' + str(j + 1)
                regression_dict[varname] = np.array(data.iloc[obj_index, :])
                regression_str += varname + '+'
            # 去掉最后一个加号
            regression_str = regression_str[:-1]
            regression_result = sm.formula.ols(regression_str, regression_dict).fit()
            # 获取拟合优度
            R_Squares.append(np.round(regression_result.rsquared_adj, 4))
            # 检测残差,异常值,正态检验
            Result = Regression.check(regression_result, regression_dict, regression_str)
            Normal_Results.append(Result)
    return R_Squares, Normal_Results



